# -*- coding: utf-8 -*-
from __future__ import annotations
import traceback
from ..utils.logger import get_logger
from ..utils.schduler import create_loop_thread, stop_loop_thread
from typing import (
    Optional,
    Union,
    List,
    Tuple,
    Callable,
    Any,
    TypeVar,
)
from .stream import TraceStream
from dataclasses import is_dataclass
import weakref
import asyncio

T = TypeVar('DataClass')
class BasePipeline():
    """
    pipeline/base.py 是pipeline的抽象基类, 定义了Pipeline的基本结构和方法。
    其主要用于
    1. streamer的全局配置和初始化。
    2. pipeline为单位的正反序列化, 方便多进程或异步框架的使用。(例如Ray, Fastapi, etc.)
    3. 采用弱引用模块, 由于Pipeline生命周期可能长短不一, 所以尽可能不让Pipeline阻碍内存回收
    """

    @classmethod
    def from_config(cls, root_cfg: T) -> BasePipeline:
        """
        从config:dataclass中提取pipeline的配置
        该方法需要被实现, 以便从配置对象中提取pipeline的配置, 并实例化pipeline
        pipeline的参数由两部分构成, 每个模块或本pipeline的dataclass, 以及自动提取环境变量
        默认情况下, 以该方法作为序列化的入口, 仅在继承该方法的情况下才可以将pipeline正反序列化
        """
        instance = cls()
        setattr(instance, "cfg", root_cfg)
        raise NotImplementedError("Pipeline must implement `from_config` method")

    @property
    def cfg(self) -> T:
        cfg = getattr(self, "_cfg", None)
        assert is_dataclass(cfg), "Pipeline config is not set"
        return cfg

    @cfg.setter
    def cfg(self, cfg: T):
        assert is_dataclass(cfg), "Pipeline config is not a dataclass"
        self._cfg = cfg

    def __reduce__(self):
        """
        用于pickle序列化, 需要返回一个可调用对象和参数
        参数需要可被反序列化, 例如dataclass或dict
        """
        return self.__class__.from_config, (self.cfg,)

    def __repr__(self):
        return f"{self.__class__.__name__}:{id(self):x}"

    @property
    def logger(self):
        logger = getattr(self, "_logger", None)
        if logger is None:
            self._logger = logger = get_logger(f"{self}")  # __repr__
        return logger
    
    @property
    def seq(self) -> List[weakref.ref[TraceStream]]:
        if seq := self._seq:
            return seq
        else:
            raise ValueError("Pipeline is not initialized")

    def __init__(
        self,
        *seq: Tuple[Union[BasePipeline, TraceStream]],
    ):
        """
        init仅把实例化后的streamer以Order的方式保存起来, 并按顺序link
        """
        self._seq = []
        for m in seq:
            if isinstance(m, TraceStream):
              self._seq.append(weakref.ref(m))
            elif isinstance(m, BasePipeline):
              self._seq.extend(m._seq)
            else:
              raise TypeError(f"Pipeline only accept streamer or pipeline, but got {type(m)}")   

    def __getitem__(self, index: Union[slice, int]) -> TraceStream:
        """
        通过索引获取streamer
        """
        if index >= len(self._seq):
            raise IndexError(f"Pipeline index out of range, but got {index}")
        return self._seq[index]()

    def __len__(self) -> int:
        return len(self.seq)

    def apply(self, func: Callable[[TraceStream], Any]):
        """
        对每个streamer应用函数, 返回结果列表
        比如用于apply(lambda x:x.logger.addHandler(...))之类
        """
        if not callable(func):
            raise TypeError(
                f"Pipeline apply only accept callable, but got {type(func)}"
            )
        results = []
        for i, module in enumerate(self.seq):
            module = module()
            if not isinstance(module, TraceStream):
                raise TypeError(
                    f"Pipeline apply only accept streamer, but got {type(module)}"
                )
            try:
                result = func(module)
            except Exception as e:
                raise RuntimeError(
                    f"Pipeline apply failed on {module.__class__.__name__}: {e}"
                )
            results.append(result)
        return results

    @property
    def top(self) -> TraceStream:
        return self.seq[0]()

    @property
    def bottom(self) -> TraceStream:
        return self.seq[-1]()
    
    def link(self, stream: TraceStream):
        self.bottom.link_to(stream)
        return BasePipeline(self, stream)

    async def __call__(
        self,
        up_module: Optional[TraceStream] = None,
        dn_module: Optional[TraceStream] = None,
    ):
        """
        异步执行函数, 返回该pipeline的协程对象
        """
        # 获取当前事件循环
        current_loop = asyncio.get_event_loop()
        if not isinstance(up_module, Optional[TraceStream]):
            raise TypeError(f"Pipeline only accept streamer, but got {type(up_module)}")
        if not isinstance(dn_module, Optional[TraceStream]):
            raise TypeError(f"Pipeline only accept streamer, but got {type(dn_module)}")
        # 集合modules
        module_instance = [m() for m in self.seq]
        if up_module is not None:
            up_module.link_to(self.top)
            module_instance.insert(0, up_module)
        if dn_module is not None:
            self.bottom.link_to(dn_module)
            module_instance.append(dn_module)        

        # 启动PorterLoop Thread闭包
        async def run_bg_loop(loop:asyncio.AbstractEventLoop):
            asyncio.set_event_loop(loop)
            _tasks = [loop.create_task(module.run_porter(loop)) for module in module_instance]
            asyncio.gather(*_tasks)
        bg_loop, bg_thread = create_loop_thread(run_bg_loop)
        
        # 在当前事件循环中启动main
        result = None
        try:
            async with asyncio.TaskGroup() as tg:
                _tasks = [tg.create_task(module.run_main(current_loop)) for module in module_instance]
        except KeyboardInterrupt:
            self.logger.warning("interrupted by user")
        except BaseExceptionGroup as e:
            for e in e.exceptions:
                if isinstance(e, StopAsyncIteration):
                    self.logger.info(f"Received StopAsyncIteration from: {e.args[0]}")
                else: # 其他未知异常抛错
                    self.logger.error(f"{''.join(traceback.format_exception(e))}")
        finally:
            # 停止后台循环
            stop_loop_thread(bg_loop, bg_thread)
            # 完成后解绑
            if up_module:
                up_module.unlink(self.top)
            if dn_module:
                result = dn_module._buffer
                self.bottom.unlink(dn_module)
            else:
                result = self.bottom._buffer
        return result
